import cv2
from ultralytics import YOLO
# 加载预训练的 YOLOv11 模型
model = YOLO("yolo11m.pt")
# 打开视频文件
video_path = r"D:\projects\dataset\JAAD\JAAD_clips\video_0109.mp4"    #video_0105.mp4
cap = cv2.VideoCapture(video_path)

# 循环读取视频帧
while cap.isOpened():
    # 从视频中读取一帧
    success, frame = cap.read()
    if success:
        results = model.predict(frame, classes=[0, 2])
        # 在帧上可视化结果
        annotated_frame = results[0].plot(line_width=2)
        # 显示带注释的帧
        cv2.imshow("YOLOv11 detection", annotated_frame)
        # 如果按下 'q' 键，退出循环
        if cv2.waitKey(20) & 0xFF == ord("q"):
            break
    else:
        # 如果视频结束，退出循环
        break
# 释放视频捕获对象并关闭显示窗口
cap.release()
cv2.destroyAllWindows()
